
import chromadb
from chromadb.config import Settings
import json
from openai import OpenAI
from dotenv import load_dotenv
import os
# 人类    鸟类
# 属性    方法

class MyVectorDBConnector:
    def __init__(self):
        # 创建数据库的链接
        self.db = chromadb.Client()
        # 创建数据库
        self.collection = self.db.get_or_create_collection(name="demo")

    def get_embeddings(self, texts, model="text-embedding-v3"):
        '''封装 qwen 的 Embedding 模型接口'''
        # print('texts', texts)
        data = client.embeddings.create(input=texts, model=model).data
        return [x.embedding for x in data]

    def add_documents(self, instructions, outputs):
        '''向 collection 中添加文档与向量'''
        # get_embeddings(instructions)
        # 将数据向量化
        embeddings = self.get_embeddings(instructions)

        # 把向量化的数据和原文存入向量数据库
        self.collection.add(
            embeddings=embeddings,  # 每个文档的向量
            documents=outputs,  # 文档的原文
            ids=[f"id{i}" for i in range(len(outputs))]  # 每个文档的 id
        )

        # print(self.collection.count())

    def search(self, query):
        '''检索向量数据库'''
        # 把我们查询的问题向量化, 在chroma当中进行查询
        results = self.collection.query(
            query_embeddings=self.get_embeddings([query]),
            n_results=2,
        )
        return results




if __name__ == '__main__':
    load_dotenv()
    client = OpenAI(api_key=os.getenv("api_key"), base_url=os.getenv("base_url"))
    # 读取文件
    with open('train_zh.json', 'r', encoding='utf-8') as f:
        data = [json.loads(line) for line in f]
    # print(data)
    # print(data[0:100])

    # 获取前10条的问题和输出
    instructions = [entry['instruction'] for entry in data[0:10]]
    outputs = [entry['output'] for entry in data[0:10]]
    # 创建一个向量数据库对象
    vector_db = MyVectorDBConnector()


    # 将问题和  回答添加到向量数据库中
    vector_db.add_documents(instructions, outputs)

    ueser_query = '白癜风'
    result = vector_db.search(ueser_query)
    print(result)







